# -*- coding: utf-8 -*-
"""
@date: 2021/6/21 17:28
@file: Lambda_use.py
@author: lilong
@desc: keras中的tensorflow使用
"""

"""
keras.layers.core.Lambda(function, output_shape=None, mask=None, arguments=None)
function：要实现的函数，该函数仅接受一个变量，即神经网络上一层的输出
mask: 掩膜
arguments：可选，是字典格式，用来传参


Lambda表达式： 用一行代码去表示一个函数;
如果某一层需要通过一个函数去变换数据，那利用keras.layers.Lambda()这个函数单独把这一步数据操作命为单独的一Lambda层。
"""


'''
def slice(x, index):

    # index是参数
    return x[:, :, index]

# 通过字典将参数index = 0传递进去
x1 = Lambda(slice, output_shape=(4, 1), arguments={'index': 0})(a)

# 通过字典将参数index = 1 传递进去
x2 = Lambda(slice, output_shape=(4, 1), arguments={'index': 1})(a)

参考：https://blog.csdn.net/weixin_43935696/article/details/112094223
'''

import numpy as np
from keras import backend as K
from keras.layers import Lambda
from keras.models import Input, Model


def lambda_1():

    # # 第一步 定义模型
    # 初始化两个输入形参
    a = Input(shape=(2,))
    b = Input(shape=(2,))

    # 定义lambda要执行的函数
    def minus(inputs):
        x, y = inputs
        return (x + y)

    # 测试
    c = [a, b]

    # 使用lambda表达式，对函数进行传参
    minus_layer = Lambda(minus, name='minus')([a, b])
    model = Model(inputs=[a, b], outputs=[minus_layer])

    # # 第二步 测试模型
    # 随便定义的两个数组
    v0 = np.array([5, 2])
    v1 = np.array([8, 4])

    # 转成1*2的矩阵后测试模型
    print(model.predict([v0.reshape(1, 2), v1.reshape(1, 2)]))


def lambda_2():
    from keras.models import Sequential
    from keras.layers import Dense, Activation, Reshape
    from keras.layers import Concatenate
    from keras.utils.vis_utils import plot_model
    from keras.layers import Input, Lambda
    from keras.models import Model

    def slice(x, index):
        return x[:, :, index]

    a = Input(shape=(4, 2))
    x1 = Lambda(slice, output_shape=(4, 1), arguments={'index': 0})(a)  # 输出维度为(4, 1)
    x2 = Lambda(slice, output_shape=(4, 1), arguments={'index': 1})(a)
    x1 = Reshape((4, 1, 1))(x1)
    x2 = Reshape((4, 1, 1))(x2)
    output = Concatenate(axis=3)([x1, x2])
    model = Model(a, output)
    plot_model(model, to_file='lambda.png', show_shapes=True)

    # 测试
    x_test = np.array([[[1, 2], [2, 3], [3, 4], [4, 5]]])
    out_test = model.predict(x_test)
    print(out_test)


def lambda_3():
    from keras.models import Sequential
    from keras.layers import Dense, Activation, Reshape
    from keras.layers import Concatenate
    from keras.utils.vis_utils import plot_model
    from keras.layers import Input, Lambda
    from keras.models import Model

    def to_one_hot(x):
        """输出一个词表大小的向量，来标记该词是否在文章出现过"""
        x, x_mask = x
        x = K.cast(x, 'int32')
        x = K.one_hot(x, 100)  # 固定词表+占位符
        x = K.sum(x_mask * x, 1, keepdims=True)
        x = K.cast(K.greater(x, 0.5), 'float32')
        return x

    x_in = Input(shape=(None,))
    y_in = Input(shape=(None,))
    x, y = x_in, y_in

    # 掩码
    x_mask = Lambda(lambda x: K.cast(K.greater(K.expand_dims(x, 2), 0), 'float32'))(x)
    x_one_hot = Lambda(to_one_hot)([x, x_mask])

    model = Model(x_in, x_one_hot)

    # 测试
    x_test = np.array(
        [[[1, 2],
          [2, 3],
          [3, 4],
          [4, 5]]]
    )
    x_test = np.load('example.npz', allow_pickle=True)['x']
    out_test = model.predict(x_test)
    print(out_test)


lambda_3()
